{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import lasagne\n",
    "from lasagne.layers import *\n",
    "import time\n",
    "import sys\n",
    "from matplotlib import pyplot as plt\n",
    "import json\n",
    "import pickle\n",
    "%matplotlib inline\n",
    "from matplotlib import patches\n",
    "import logging\n",
    "from os.path import join, exists\n",
    "from os import mkdir, makedirs\n",
    "from collections import Counter\n",
    "from notebooks.nbfinder import NotebookFinder\n",
    "sys.meta_path.append(NotebookFinder())\n",
    "from notebooks.print_n_plot import *\n",
    "from notebooks.helper_fxns import *\n",
    "from sklearn.metrics import average_precision_score\n",
    "import copy\n",
    "\n",
    "class TrainVal(object):\n",
    "    def __init__(self, iterator, kwargs, fns, networks):\n",
    "        self.metrics = {}\n",
    "        self.kwargs = kwargs\n",
    "        self.iterator = iterator\n",
    "        self.fns = fns\n",
    "        self.xdim=768\n",
    "        self.ydim = 1152\n",
    "        self.epoch = 0\n",
    "        self.start_time = 0\n",
    "        self.seed = 5\n",
    "        self.classes = [ \"TD\", \"TC\", \"ETC\", \"AR\"]\n",
    "        self.max_ims = self.kwargs['num_ims_to_plot']\n",
    "        self.networks = networks\n",
    "        self.print_network(networks)\n",
    "        it_list= [\"variables\",'batch_size',\"data_dir\", \"metadata_dir\", \"shuffle\",\"num_classes\",\n",
    "                    \"labels_only\",\"time_chunks_per_example\"]\n",
    "        \n",
    "        self.std_it_kwargs = {k:kwargs[k] for k in it_list} \n",
    "        self.it_kwargs = {k:copy.deepcopy(self.std_it_kwargs) for k in [\"tr\", \"val\", \"test\"]}\n",
    "        for k,v in self.it_kwargs.iteritems():\n",
    "            self.it_kwargs[k].update({'days':kwargs['num_'+ k+ '_days'], 'years':kwargs[k +'_years'], \"seed\": 5})\n",
    "        \n",
    "        \n",
    "    \n",
    "        \n",
    "        self.plotter = Plotter(self.kwargs, self.fns, self.max_ims, self.iterator)\n",
    "    def print_network(self, networks):\n",
    "        yolo, ae = networks['yolo'], networks['ae']\n",
    "        self._print_network(yolo)\n",
    "        if self.kwargs['lambda_ae'] != 0:\n",
    "            \n",
    "            self._print_network(ae)\n",
    "            \n",
    "            \n",
    "    def _print_network(self, network):\n",
    "        self.kwargs['logger'].info(\"\\n\")\n",
    "        for layer in get_all_layers(network):\n",
    "            self.kwargs['logger'].info(str(layer) +' : ' + str(layer.output_shape))\n",
    "        self.kwargs['logger'].info(str(count_params(layer)))\n",
    "        self.kwargs['logger'].info(\"\\n\")\n",
    "    \n",
    "    def do_one_epoch(self):\n",
    "        self._do_one_epoch(type_=\"tr\")\n",
    "        self._do_one_epoch(type_=\"val\")\n",
    "        self.print_results()\n",
    "        self.plotter.plot_learn_curve(self.metrics)\n",
    "        self.epoch += 1\n",
    "    def _do_one_epoch(self, type_=\"tr\"):\n",
    "        print \"beginning epoch %i\" % (self.epoch)\n",
    "        start_time = time.time()\n",
    "        loss_tots = {}\n",
    "        acc_tots = {}\n",
    "        batches = 0\n",
    "        it_kwargs = self.it_kwargs[type_]\n",
    "        \n",
    "     \n",
    "        t0 = time.time()\n",
    "        for x,y in self.iterator(**it_kwargs).iterate():\n",
    "            print \"total iterate time: \", time.time() -t0\n",
    "            t= time.time()\n",
    "            loss_dict, acc_dict = self.do_one_iteration(x,y,type_)\n",
    "            print \"actual iteration time: \", time.time() - t\n",
    "            for k in loss_dict.keys():\n",
    "                key = type_ + \"_\" + k\n",
    "                loss_tots = add_as_running_total(key,loss_dict[k], loss_tots)\n",
    "            \n",
    "                        \n",
    "            for k in acc_dict.keys():\n",
    "                key = type_ + \"_\" + k\n",
    "                acc_tots = add_as_extension(key,acc_dict[k], acc_tots)\n",
    "                \n",
    "            \n",
    "            if not self.kwargs[\"no_plots\"]:\n",
    "                self.plotter.do_plots(x, y, type_, batches,self.epoch)\n",
    "                           \n",
    "   \n",
    "            batches += 1\n",
    "            t0 = time.time()\n",
    "        t= time.time()\n",
    "        self.postprocess_epoch(type_, loss_tots,acc_tots, start_time, batches)\n",
    "        print \"postprocess time: \", time.time() - t\n",
    "      \n",
    "        if type_ == \"val\":\n",
    "            self.save_weights()\n",
    "            \n",
    "    def do_one_iteration(self, x,y,type_):\n",
    "        loss_dict = self.fns[type_](x,y)\n",
    "            \n",
    "        acc_dict = self.fns[\"acc\"](x,y, iou_thresh=self.kwargs['iou_thresh'],\n",
    "                                            conf_thresh=self.kwargs['conf_thresh'])            \n",
    "           \n",
    "        return loss_dict, acc_dict #, pred_boxes, gt_boxes\n",
    "        \n",
    "\n",
    "    def postprocess_epoch(self, type_, loss_tots,acc_tots, start_time, batches):\n",
    "        assert batches > 0\n",
    "                \n",
    "        loss_tots = {k: v / float(batches) for k,v in loss_tots.iteritems()}\n",
    "        acc_tots_by_class = {}\n",
    "        for k in range(self.kwargs[\"num_classes\"]):\n",
    "            gt_ans = acc_tots[type_ + \"_gt_\" + str(k)]\n",
    "            pred_confs = acc_tots[type_ + \"_pred_\" + str(k)]\n",
    "            if len(pred_confs) > 0:\n",
    "                ap = average_precision_score(gt_ans, pred_confs)\n",
    "                if np.isnan(ap):\n",
    "                    ap = 0.\n",
    "\n",
    "            # there are no predictions and no actual labels, so undefined precision\n",
    "            else:\n",
    "                ap = np.nan\n",
    "            acc_tots_by_class[type_ + \"_\" + self.classes[k] + \"_ap\"] = ap\n",
    "\n",
    "        \n",
    "        #don't penalize for when there are no ground truth of a class\n",
    "        mAP_list = [a for a in acc_tots_by_class.values() if not np.isnan(a)]\n",
    "        mAP = np.mean(mAP_list)\n",
    "        for k,v in loss_tots.iteritems():\n",
    "            if not isinstance(v, float):\n",
    "                print \"adding \", k, \" equal to \", v\n",
    "            self.metrics = add_as_appension(k,v,self.metrics)\n",
    "        \n",
    "        for k,v in acc_tots_by_class.iteritems():\n",
    "            if not isinstance(v, float):\n",
    "                print \"adding \", k, \" equal to \", v\n",
    "            self.metrics = add_as_appension(k,v,self.metrics)\n",
    "        \n",
    "        if not isinstance(mAP, float):\n",
    "                print \"adding \", \"map\", \" equal to \", mAP\n",
    "        self.metrics = add_as_appension(type_ + \"_mAP\", mAP,self.metrics)\n",
    "        \n",
    "        time_key = type_ + \"_time\"\n",
    "        self.metrics = add_as_appension(time_key,time.time() - start_time, self.metrics)\n",
    "        \n",
    "        \n",
    "\n",
    "    \n",
    "    def save_weights(self):\n",
    "        #print self.metrics\n",
    "        max_metrics = [\"val_mAP\"]\n",
    "        min_metrics = [\"val_loss\"]\n",
    "        for k in max_metrics:\n",
    "            if len(self.metrics[k]) > 1:\n",
    "                if self.metrics[k][-1] > max(self.metrics[k][:-1]):\n",
    "                    self._save_weights(\"yolo\", \"best_\" + k)\n",
    "        \n",
    "        \n",
    "            else:\n",
    "                self._save_weights(\"yolo\", \"best_\" + k)\n",
    "        for k in min_metrics:\n",
    "            if len(self.metrics[k]) > 1:\n",
    "                if self.metrics[k][-1] < min(self.metrics[k][:-1]):\n",
    "                    self._save_weights(\"yolo\", \"best_\" + k)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        self._save_weights(\"yolo\", \"cur\")\n",
    "        self._save_weights(\"ae\", \"cur\")\n",
    "        \n",
    "    def test(self):\n",
    "        self._do_one_epoch(type_=\"test\")\n",
    "        self.print_results(type_=\"test\")\n",
    "    def val(self):\n",
    "        self._do_one_epoch(type_=\"val\")\n",
    "        self.print_results(type_=\"val\")\n",
    "    \n",
    "    def train(self):\n",
    "        for epoch in range(self.kwargs['epochs']):\n",
    "            self.do_one_epoch()\n",
    "        \n",
    "        \n",
    "\n",
    "        \n",
    "        \n",
    "    def _save_weights(self,name,suffix=\"\"):\n",
    "        params = get_all_param_values(self.networks[name])\n",
    "        model_dir = join(self.kwargs['save_path'], \"models\")\n",
    "        makedir_if_not_there(model_dir)\n",
    "        pickle.dump(params,open(join(model_dir, name + \"_\" + suffix + \".pkl\"), \"w\"))\n",
    "    \n",
    "    def print_results(self,type_=None):\n",
    "        self.kwargs['logger'].info(\"Epoch {} of {}\".format(self.epoch + 1, self.kwargs['epochs']))\n",
    "        for typ in [\"tr\", \"val\", \"test\"]:\n",
    "            if type_ is not None:\n",
    "                if typ != type_:\n",
    "                    continue\n",
    "            else:\n",
    "                if typ == \"test\":\n",
    "                    continue\n",
    "            self.kwargs['logger'].info(\"\\t {} took {:.3f}s\".format(typ, self.metrics[typ + \"_time\"][-1]))\n",
    "            for k,v in self.metrics.iteritems():\n",
    "                if typ in k[:5] and \"time\" not in k:\n",
    "                  \n",
    "                    if \"acc\" in k:\n",
    "                        self.kwargs['logger'].info(\"\\t\\t\" + k + \":\\t\\t{:.4f} %\".format(v[-1] * 100))\n",
    "                    else:\n",
    "                        self.kwargs['logger'].info(\"\\t\\t\" + k + \":\\t\\t{:.4f}\".format(v[-1]))\n",
    "\n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "    \n",
    "    \n",
    "# def get_fmaps(iterator, kwargs, networks, fns):\n",
    "#     tv = TrainVal(iterator,kwargs, fns, networks)\n",
    "#     tv.get_encoder_fmaps()\n",
    "\n",
    "# def get_ims(iterator, kwargs, networks, fns):\n",
    "#     tv = TrainVal(iterator,kwargs, fns, networks)\n",
    "#     tv.postproc_ims()\n",
    "    \n",
    "    \n",
    "    \n",
    "         "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
